{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "header",
   "metadata": {},
   "source": [
    "# Kubeflow Trainer: Container Backend (Single-Node) Training\n",
    "\n",
    "This notebook demonstrates how to run single-node training using the **Container Backend** with Docker or Podman.\n",
    "\n",
    "## Container Backend\n",
    "\n",
    "- **Container Runtime**: Docker or Podman required\n",
    "- **Use Case**: Testing container workflows, simulating production environments\n",
    "- **Prerequisites**: Python 3.9+ and Docker Desktop/Engine OR Podman\n",
    "\n",
    "This example trains a CNN on the classic [MNIST](http://yann.lecun.com/exdb/mnist/) handwritten digit dataset using PyTorch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "install",
   "metadata": {},
   "source": [
    "## Install the Kubeflow SDK\n",
    "\n",
    "You need to install the Kubeflow SDK with container backend support:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pip-install",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to install\n",
    "# %pip install -U kubeflow[docker]  # For Docker\n",
    "# %pip install -U kubeflow[podman]  # For Podman"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "training-function",
   "metadata": {},
   "source": [
    "## Define the Training Function\n",
    "\n",
    "The first step is to create a function to train CNN model using MNIST data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "train-function",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_mnist():\n",
    "    import os\n",
    "    import torch\n",
    "    import torch.nn.functional as F\n",
    "    from torch import nn, optim\n",
    "    import torch.distributed as dist\n",
    "    from torch.utils.data import DataLoader, DistributedSampler\n",
    "    from torchvision import datasets, transforms\n",
    "\n",
    "    # Define the PyTorch CNN model to be trained\n",
    "    class Net(nn.Module):\n",
    "        def __init__(self):\n",
    "            super(Net, self).__init__()\n",
    "            self.conv1 = nn.Conv2d(1, 20, 5, 1)\n",
    "            self.conv2 = nn.Conv2d(20, 50, 5, 1)\n",
    "            self.fc1 = nn.Linear(4 * 4 * 50, 500)\n",
    "            self.fc2 = nn.Linear(500, 10)\n",
    "\n",
    "        def forward(self, x):\n",
    "            x = F.relu(self.conv1(x))\n",
    "            x = F.max_pool2d(x, 2, 2)\n",
    "            x = F.relu(self.conv2(x))\n",
    "            x = F.max_pool2d(x, 2, 2)\n",
    "            x = x.view(-1, 4 * 4 * 50)\n",
    "            x = F.relu(self.fc1(x))\n",
    "            x = self.fc2(x)\n",
    "            return F.log_softmax(x, dim=1)\n",
    "\n",
    "    # Use NCCL if a GPU is available, otherwise use Gloo as communication backend.\n",
    "    device, backend = (\"cuda\", \"nccl\") if torch.cuda.is_available() else (\"cpu\", \"gloo\")\n",
    "    print(f\"Using Device: {device}, Backend: {backend}\")\n",
    "\n",
    "    # Setup PyTorch distributed.\n",
    "    local_rank = int(os.getenv(\"LOCAL_RANK\", 0))\n",
    "    dist.init_process_group(backend=backend)\n",
    "    print(\n",
    "        \"Distributed Training for WORLD_SIZE: {}, RANK: {}, LOCAL_RANK: {}\".format(\n",
    "            dist.get_world_size(),\n",
    "            dist.get_rank(),\n",
    "            local_rank,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    # Create the model and load it into the device.\n",
    "    device = torch.device(f\"{device}:{local_rank}\")\n",
    "    model = nn.parallel.DistributedDataParallel(Net().to(device))\n",
    "    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)\n",
    "\n",
    "    # Download FashionMNIST dataset only on local_rank=0 process.\n",
    "    if local_rank == 0:\n",
    "        dataset = datasets.FashionMNIST(\n",
    "            \"./data\",\n",
    "            train=True,\n",
    "            download=True,\n",
    "            transform=transforms.Compose([\n",
    "                transforms.ToTensor(),\n",
    "                transforms.Normalize((0.1307,), (0.3081,)),\n",
    "            ]),\n",
    "        )\n",
    "    dist.barrier()\n",
    "    dataset = datasets.FashionMNIST(\n",
    "        \"./data\",\n",
    "        train=True,\n",
    "        download=False,\n",
    "        transform=transforms.Compose([transforms.ToTensor()]),\n",
    "    )\n",
    "\n",
    "    # Shard the dataset accross workers.\n",
    "    train_loader = DataLoader(\n",
    "        dataset,\n",
    "        batch_size=100,\n",
    "        sampler=DistributedSampler(dataset)\n",
    "    )\n",
    "\n",
    "    dist.barrier()\n",
    "    for epoch in range(1, 3):\n",
    "        model.train()\n",
    "        \n",
    "        # Iterate over mini-batches from the training set\n",
    "        for batch_idx, (data, target) in enumerate(train_loader):\n",
    "            # Forward pass\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            outputs = model(data)\n",
    "            loss = F.nll_loss(outputs, target)\n",
    "            # Backward pass\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            if batch_idx % 100 == 0:\n",
    "                print(\n",
    "                    \"Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}\".format(\n",
    "                        epoch,\n",
    "                        batch_idx * len(data),\n",
    "                        len(train_loader.dataset),\n",
    "                        100.0 * batch_idx / len(train_loader),\n",
    "                        loss.item(),\n",
    "                    )\n",
    "                )\n",
    "\n",
    "    # Wait for the distributed training to complete\n",
    "    dist.barrier()\n",
    "    if dist.get_rank() == 0:\n",
    "        print(\"Training is finished\")\n",
    "        torch.save(model.state_dict(), \"mnist_cnn.pt\")\n",
    "\n",
    "    # Finally clean up PyTorch distributed\n",
    "    dist.destroy_process_group()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "configure-backend",
   "metadata": {},
   "source": [
    "## Configure Container Backend\n",
    "\n",
    "The container backend automatically detects and uses either Docker or Podman:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "backend-config",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kubeflow.trainer import TrainerClient, ContainerBackendConfig\n",
    "import os\n",
    "\n",
    "# Auto-detects Docker or Podman\n",
    "backend_config = ContainerBackendConfig()\n",
    "\n",
    "# Optional: Force specific runtime\n",
    "# backend_config = ContainerBackendConfig(runtime=\"docker\")  # Force Docker\n",
    "# backend_config = ContainerBackendConfig(runtime=\"podman\")  # Force Podman\n",
    "\n",
    "# Optional: For Colima on macOS\n",
    "# backend_config = ContainerBackendConfig(\n",
    "#     container_host=f\"unix://{os.path.expanduser('~')}/.colima/default/docker.sock\"\n",
    "# )\n",
    "\n",
    "# Optional: For Podman Machine on macOS\n",
    "# backend_config = ContainerBackendConfig(\n",
    "#     runtime=\"podman\",\n",
    "#     container_host=\"unix:///run/user/1000/podman/podman.sock\"\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "init-client",
   "metadata": {},
   "source": [
    "## Initialize Client\n",
    "\n",
    "Initialize the TrainerClient with the Container Backend:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "create-client",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = TrainerClient(backend_config=backend_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "list-runtimes",
   "metadata": {},
   "source": [
    "## List the Training Runtimes\n",
    "\n",
    "You can get the list of available Training Runtimes to start your TrainJob."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "get-runtimes",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Runtime(name='deepspeed-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='deepspeed', image='ghcr.io/kubeflow/trainer/deepspeed-runtime', num_nodes=1, device='Unknown', device_count='Unknown'), pretrained_model=None)\n",
      "Runtime(name='mlx-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='mlx', image='ghcr.io/kubeflow/trainer/mlx-runtime', num_nodes=1, device='Unknown', device_count='Unknown'), pretrained_model=None)\n",
      "Runtime(name='torch-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='torch', image='pytorch/pytorch:2.7.1-cuda12.8-cudnn9-runtime', num_nodes=1, device='Unknown', device_count='Unknown'), pretrained_model=None)\n"
     ]
    }
   ],
   "source": [
    "for runtime in client.list_runtimes():\n",
    "    print(runtime)\n",
    "    if runtime.name == \"torch-distributed\":\n",
    "        torch_runtime = runtime"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "start-training",
   "metadata": {},
   "source": [
    "## Run the TrainJob\n",
    "\n",
    "Submit the training job to the Container Backend (single container):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "train-job",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kubeflow.trainer import CustomTrainer\n",
    "\n",
    "job_name = client.train(\n",
    "    trainer=CustomTrainer(\n",
    "        func=train_mnist,\n",
    "        packages_to_install=[\"torchvision\"],\n",
    "        num_nodes=2,\n",
    "    ),\n",
    "    runtime=torch_runtime,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "check-status",
   "metadata": {},
   "source": [
    "## Check the TrainJob Status\n",
    "\n",
    "You can check the status of the TrainJob that's created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "job-status",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Job: p6072ef2ca48, Status: Running\n"
     ]
    }
   ],
   "source": [
    "job = client.get_job(job_name)\n",
    "print(\"Job: {}, Status: {}\".format(job.name, job.status))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stream-logs",
   "metadata": {},
   "source": [
    "## Watch the TrainJob Logs\n",
    "\n",
    "We can use the `get_job_logs()` API to get the TrainJob logs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "logs",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torchvision in /opt/conda/lib/python3.11/site-packages (0.22.1+cu128)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.11/site-packages (from torchvision) (2.2.6)\n",
      "Requirement already satisfied: torch==2.7.1 in /opt/conda/lib/python3.11/site-packages (from torchvision) (2.7.1+cu128)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /opt/conda/lib/python3.11/site-packages (from torchvision) (11.0.0)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (3.18.0)\n",
      "Requirement already satisfied: typing-extensions>=4.10.0 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (4.14.0)\n",
      "Requirement already satisfied: sympy>=1.13.3 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (1.14.0)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (3.5)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (3.1.6)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (2025.5.1)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.61 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.61)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.57 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.57)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.57 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.57)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.7.1.26 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (9.7.1.26)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.8.3.14 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.3.14)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.3.3.41 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (11.3.3.41)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.9.55 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (10.3.9.55)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.7.2.55 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (11.7.2.55)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.5.7.53 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.5.7.53)\n",
      "Requirement already satisfied: nvidia-cusparselt-cu12==0.6.3 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (0.6.3)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.26.2 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (2.26.2)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.8.55 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.55)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.61 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.61)\n",
      "Requirement already satisfied: nvidia-cufile-cu12==1.13.0.11 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (1.13.0.11)\n",
      "Requirement already satisfied: triton==3.3.1 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (3.3.1)\n",
      "Requirement already satisfied: setuptools>=40.8.0 in /opt/conda/lib/python3.11/site-packages (from triton==3.3.1->torch==2.7.1->torchvision) (75.8.2)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /opt/conda/lib/python3.11/site-packages (from sympy>=1.13.3->torch==2.7.1->torchvision) (1.3.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.11/site-packages (from jinja2->torch==2.7.1->torchvision) (3.0.2)\n",
      "WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\n",
      "Using Device: cpu, Backend: gloo\n",
      "Distributed Training for WORLD_SIZE: 2, RANK: 0, LOCAL_RANK: 0\n",
      "100%|██████████| 26.4M/26.4M [00:00<00:00, 50.3MB/s]\n",
      "100%|██████████| 29.5k/29.5k [00:00<00:00, 2.97MB/s]\n",
      "100%|██████████| 4.42M/4.42M [00:00<00:00, 43.0MB/s]\n",
      "100%|██████████| 5.15k/5.15k [00:00<00:00, 47.2MB/s]\n",
      "Train Epoch: 1 [0/60000 (0%)]\tLoss: 2.315568\n",
      "Train Epoch: 1 [10000/60000 (33%)]\tLoss: 0.734196\n",
      "Train Epoch: 1 [20000/60000 (67%)]\tLoss: 0.756162\n",
      "Train Epoch: 2 [0/60000 (0%)]\tLoss: 0.722734\n",
      "Train Epoch: 2 [10000/60000 (33%)]\tLoss: 0.428438\n",
      "Train Epoch: 2 [20000/60000 (67%)]\tLoss: 0.625209\n",
      "Training is finished\n",
      "Requirement already satisfied: torchvision in /opt/conda/lib/python3.11/site-packages (0.22.1+cu128)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.11/site-packages (from torchvision) (2.2.6)\n",
      "Requirement already satisfied: torch==2.7.1 in /opt/conda/lib/python3.11/site-packages (from torchvision) (2.7.1+cu128)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /opt/conda/lib/python3.11/site-packages (from torchvision) (11.0.0)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (3.18.0)\n",
      "Requirement already satisfied: typing-extensions>=4.10.0 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (4.14.0)\n",
      "Requirement already satisfied: sympy>=1.13.3 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (1.14.0)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (3.5)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (3.1.6)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (2025.5.1)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.61 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.61)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.57 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.57)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.57 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.57)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.7.1.26 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (9.7.1.26)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.8.3.14 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.3.14)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.3.3.41 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (11.3.3.41)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.9.55 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (10.3.9.55)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.7.2.55 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (11.7.2.55)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.5.7.53 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.5.7.53)\n",
      "Requirement already satisfied: nvidia-cusparselt-cu12==0.6.3 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (0.6.3)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.26.2 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (2.26.2)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.8.55 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.55)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.61 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (12.8.61)\n",
      "Requirement already satisfied: nvidia-cufile-cu12==1.13.0.11 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (1.13.0.11)\n",
      "Requirement already satisfied: triton==3.3.1 in /opt/conda/lib/python3.11/site-packages (from torch==2.7.1->torchvision) (3.3.1)\n",
      "Requirement already satisfied: setuptools>=40.8.0 in /opt/conda/lib/python3.11/site-packages (from triton==3.3.1->torch==2.7.1->torchvision) (75.8.2)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /opt/conda/lib/python3.11/site-packages (from sympy>=1.13.3->torch==2.7.1->torchvision) (1.3.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.11/site-packages (from jinja2->torch==2.7.1->torchvision) (3.0.2)\n",
      "WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\n",
      "Waiting for master node p6072ef2ca48-node-0:29500...\n",
      "Master node is reachable\n",
      "Using Device: cpu, Backend: gloo\n",
      "Distributed Training for WORLD_SIZE: 2, RANK: 1, LOCAL_RANK: 0\n",
      "100%|██████████| 26.4M/26.4M [00:00<00:00, 51.8MB/s]\n",
      "100%|██████████| 29.5k/29.5k [00:00<00:00, 2.81MB/s]\n",
      "100%|██████████| 4.42M/4.42M [00:00<00:00, 25.2MB/s]\n",
      "100%|██████████| 5.15k/5.15k [00:00<00:00, 13.3MB/s]\n",
      "Train Epoch: 1 [0/60000 (0%)]\tLoss: 2.322156\n",
      "Train Epoch: 1 [10000/60000 (33%)]\tLoss: 0.986901\n",
      "Train Epoch: 1 [20000/60000 (67%)]\tLoss: 0.654047\n",
      "Train Epoch: 2 [0/60000 (0%)]\tLoss: 0.686302\n",
      "Train Epoch: 2 [10000/60000 (33%)]\tLoss: 0.607211\n",
      "Train Epoch: 2 [20000/60000 (67%)]\tLoss: 0.552617\n"
     ]
    }
   ],
   "source": [
    "for logline in client.get_job_logs(job_name, follow=True):\n",
    "    print(logline, end='')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cleanup",
   "metadata": {},
   "source": [
    "## Delete the TrainJob\n",
    "\n",
    "When the TrainJob is finished, you can delete the resource."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "delete",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.delete_job(job_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
